import json

import pandas as pd
import requests
from more_itertools import chunked
from tqdm import tqdm

dataset_id = "ef72cca0-1fb0-40e5-8e63-40c419ebe9ad"
document_id = "74e29d2f-b191-4ed3-a33b-3d3c00f59498"
df = pd.read_excel("./.vscode/恋爱助手/恋爱军师生成偏好.xlsx")

url = f"http://localhost:5001/v1/datasets/{dataset_id}/documents/{document_id}/segments"

headers = {
    "Authorization": "Bearer dataset-7m0rE371MEiKylNb8iDYXBny",
    "Content-Type": "application/json",
}

for batch in tqdm(chunked(df.iterrows(), 5), total=len(df) // 5):
    data = {
        "segments": [
            {
                "content": f"{row.iloc[0]}",
                "answer": json.dumps(row.to_dict(), ensure_ascii=False),
                "keywords": [row.iloc[0]],
            }
            for _, row in batch
        ]
    }
    response = requests.post(
        url, headers=headers, data=json.dumps(data, ensure_ascii=False)
    )
    print(response.text)
